import os
import numpy as np
import tensorflow as tf
from prepocess.utils import get_file, get_batch_tensor

from nets.cnn_example import inference, losses, trainning, evaluation
from nets.mobilenet import mobilenet_v2

IMG_W = 224  # resize图像，太大的话训练时间久
IMG_H = 224
CLASS_NUM = 574
BATCH_SIZE = 32  # 每个batch要放多少张图片
CAPACITY = 640  # 一个队列最大多少
MAX_STEP = 5000000  # 一般大于10K
learning_rate = 0.00001  # 一般小于0.0001
restore_step = 0
ckpt = {}

# 获取批次batch
train_dir = './data_set_lit/'  # 训练样本的读入路径
# train_dir = 'D:/PycharmProjects/qiyi_com/mix_part1'
# train_dir = 'D:/resources/data_set/MNIST_MIX'
logs_train_dir = './train'  # logs存储路径
train, train_label, val, val_label, class_names_to_ids,ids_to_class_names= get_file(train_dir)
print(" train files : {} , val files {}".format(len(train),len(val)))


train_batch, train_label_batch = get_batch_tensor(train, train_label, IMG_W, IMG_H, BATCH_SIZE, CAPACITY,class_num=CLASS_NUM)
val_batch, val_label_batch = get_batch_tensor(val, val_label, IMG_W, IMG_H, BATCH_SIZE, CAPACITY,class_num=CLASS_NUM)


with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope()):
    train_logits, endpoints = mobilenet_v2.mobilenet(train_batch,num_classes=CLASS_NUM, depth_multiplier=1.4)
    train_loss = losses(train_logits, train_label_batch)
    train_op = trainning(train_loss, learning_rate)
    train_acc = evaluation(train_logits, train_label_batch, 'train_acc')
    val_acc = evaluation(train_logits, val_label_batch, 'val_acc')
    summary_op = tf.summary.merge_all()


sess = tf.Session()
train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
# 产生一个saver来存储训练好的模型

#
# exclude = ['Variable_12', 'Variable_13', 'Variable_11', 'Variable_5']
# variables_to_restore = slim.get_variables_to_restore(exclude=exclude)
#
# saver = tf.train.Saver(variables_to_restore)
saver = tf.train.Saver(max_to_keep=3)
# if os.path.exists(logs_train_dir):
#     ckpt = tf.train.get_checkpoint_state(logs_train_dir)
#     print("loading {}".format(ckpt.model_checkpoint_path))
#     saver.restore(sess, os.path.join(ckpt.model_checkpoint_path))
#     restore_step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1])

# 所有节点初始化
sess.run(tf.global_variables_initializer())
# 队列监控
coord = tf.train.Coordinator()  # 设置多线程协调器
threads = tf.train.start_queue_runners(sess=sess, coord=coord)


try:
    for step in np.arange(restore_step,MAX_STEP):
        if coord.should_stop():
            break
        _, tra_loss, tra_acc, va_acc = sess.run([train_op, train_loss, train_acc, val_acc])

        if step % 100 == 0:
            print('Step %d, train loss = %.2f, train accuracy = %.2f%% val accuracy = %.2f%% ' % (
                step, tra_loss, tra_acc * 100.0, va_acc * 100.0))
            summary_str = sess.run(summary_op)
            train_writer.add_summary(summary_str, step)
            checkpoint_path = os.path.join(logs_train_dir, 'mobilenet_v2_1.0_224.ckpt')
            saver.save(sess, checkpoint_path,global_step= step)
        '''
        # 每隔100步，保存一次训练好的模型
        if (step + 1) == MAX_STEP:
            checkpoint_path = os.path.join(logs_train_dir, 'thing.ckpt')
            saver.save(sess, checkpoint_path, global_step=step)
        '''

except tf.errors.OutOfRangeError:
    print('Done training -- epoch limit reached')

finally:
    coord.request_stop()
coord.join(threads)
sess.close()

